package com.anttribe.components.infra.lang.reflect.scanner;

import com.anttribe.components.infra.lang.reflect.scanner.exception.ClassScannerException;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;

import java.io.IOException;
import java.net.JarURLConnection;
import java.net.URL;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.Set;
import java.util.function.Predicate;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

/**
 * @author zhaoyong
 * @date 2021/11/4 0004
 */
public class JarClassScanner implements ClassScanner {

    private static final String PROTOCOL_JAR = "jar";

    private static final String PROTOCOL_FILE = "file";

    @Override
    public Set<Class<?>> scan(String basePackage, Predicate<Class<?>> predicate) {
        return doScan(basePackage, predicate);
    }

    /**
     * 执行扫描任务
     *
     * @param packageName 包名
     * @param predicate   筛选
     * @return Set<Class < ?>>
     */
    private Set<Class<?>> doScan(String packageName, Predicate<Class<?>> predicate) {
        Set<Class<?>> scanClasses = new HashSet<>();
        String scanFilePath = packageName.replace(PACKAGE_SEPARATOR, FILE_SEPARATOR);
        try {
            // 通过当前线程得到类加载器从而得到URL的枚举
            Enumeration<URL> urlEnumeration = Thread.currentThread().getContextClassLoader().getResources(scanFilePath);
            while (urlEnumeration.hasMoreElements()) {
                URL url = urlEnumeration.nextElement();
                String protocol = url.getProtocol();
                Set<Class<?>> urlScanClasses = null;
                if (PROTOCOL_JAR.equalsIgnoreCase(protocol)) {
                    urlScanClasses = doScan(packageName, url, predicate);
                } else if (PROTOCOL_FILE.equalsIgnoreCase(protocol)) {
                    // 从maven子项目中扫描
                    FileClassScanner fileClassScanner = new FileClassScanner();
                    fileClassScanner.setClasspath(url.getPath().replace(scanFilePath, ""));
                    urlScanClasses = fileClassScanner.scan(packageName, predicate);
                }
                if (!CollectionUtils.isEmpty(urlScanClasses)) {
                    scanClasses.addAll(urlScanClasses);
                }
            }
        } catch (IOException e) {
            throw new ClassScannerException(e);
        }
        return scanClasses;
    }

    /**
     * 执行扫描任务
     *
     * @param packageName 包名
     * @param url         jar文件
     * @param predicate   筛选
     * @return Set<Class < ?>>
     */
    private Set<Class<?>> doScan(String packageName, URL url, Predicate<Class<?>> predicate) {
        Set<Class<?>> scanClasses = new HashSet<>();
        try {
            // 转换为JarURLConnection
            JarURLConnection urlConnection = (JarURLConnection) url.openConnection();
            if (null == urlConnection) {
                return scanClasses;
            }

            JarFile jarFile = urlConnection.getJarFile();
            if (null == jarFile) {
                return scanClasses;
            }

            // 得到该jar文件下面的类实体
            Enumeration<JarEntry> jarEntryEnumeration = jarFile.entries();
            while (jarEntryEnumeration.hasMoreElements()) {
                JarEntry jarEntity = jarEntryEnumeration.nextElement();
                String jarEntryName = jarEntity.getName();
                // 需要过滤class文件并且在指定包名下的类
                if (!StringUtils.endsWith(jarEntryName, CLASS_SUFFIX)) {
                    continue;
                }

                String classCanonicalName = jarEntryName.replaceAll(FILE_SEPARATOR, PACKAGE_SEPARATOR);
                if (StringUtils.startsWith(classCanonicalName, packageName)) {
                    String className = classCanonicalName.substring(0, jarEntryName.lastIndexOf("."));
                    Class clazz = Class.forName(className, Boolean.FALSE, Thread.currentThread().getContextClassLoader());
                    if (predicate == null || predicate.test(clazz)) {
                        scanClasses.add(clazz);
                    }
                }
            }
        } catch (IOException | ClassNotFoundException e) {
            throw new ClassScannerException(e);
        }
        return scanClasses;
    }

}
